import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path

from sqlmodel import select, func

from markitdown_api.app.config.db import get_session_async_ctx
from markitdown_api.app.config.settings import settings
from markitdown_api.app.handler.task_handler import TaskResult
from markitdown_api.app.handler.task_handler_factory import task_handler_factory
from src.markitdown_api.app.models.task import TaskStatus, Task, find_task_by_id, save_task

logger = logging.getLogger("taskScheduleHandler")
logger.setLevel(logging.DEBUG)


class _TaskScheduleHandler:
    """
    任务调度处理器
    """

    def __init__(self, max_worker: int):
        """
        :param max_worker:  最大并发限制
        """
        logger.info(f"初始化任务调度处理器，最大工作线程：{max_worker}")
        self.max_consumers = max_worker
        self.semaphore = asyncio.Semaphore(max_worker)
        self.executor = ThreadPoolExecutor(max_workers=max_worker)
        self.closed = False
        self.event = asyncio.Event()

    async def schedule_tasks(self):
        """
        定时任务：每隔 15 秒从队列中获取任务并提交到线程池
        :return:
        """

        while True:
            if self.closed is True:
                break

            if not self.semaphore.locked():
                async with get_session_async_ctx() as session:
                    statement = select(func.count()).where(Task.status == TaskStatus.PENDING)
                    count = (await session.execute(statement)).scalar()
                    logger.info(f"查询到 {count} 个待执行任务...")

                if count > 0:
                    while True:
                        # 无许可证结束任务循环
                        if self.semaphore.locked():
                            break

                        async with get_session_async_ctx() as session:
                            statement = select(Task).where(Task.status == TaskStatus.PENDING).order_by(Task.id).limit(1)
                            task = (await session.execute(statement)).scalar_one_or_none()

                        if task is None:
                            logger.info(f"没有更多的待执行任务了，本次调度结束...")
                            break

                        # 等待信号量（确保不超过max_consumers个待处理任务）
                        await self.semaphore.acquire()

                        if self.closed is True:
                            break

                        logger.info(f"提交任务 {task.id} 导致工作线程执行...")
                        asyncio.create_task(self.markitdown_task(task.id))
                else:
                    logger.info(f"没有待执行的任务，本次调度结束...")
                    if self.closed is True:
                        break

            # 每隔 15 秒执行一次
            try:
                self.event.clear()
                await asyncio.wait_for(self.event.wait(), timeout=60)
                # 事件唤醒
            except asyncio.TimeoutError:
                # 超时唤醒
                continue

    async def trigger_schedule(self):
        """触发调度"""
        self.event.set()

    async def markitdown_task(self, task_id: int):
        """
        执行 markitdown 任务
        :param task_id: 任务id
        :return:
        """
        try:
            logger.info(f"开始执行任务 {task_id} ...")

            async with get_session_async_ctx() as session:
                task = await find_task_by_id(session, task_id)

            if not task:
                logger.warning(f"任务id {task_id} 不存在")
                return

            # 任务执行中
            async with get_session_async_ctx() as session:
                task.status = TaskStatus.EXECUTING
                task.start_execution_time = datetime.now()
                task = await save_task(session, task)

            try:
                # 将同步代码放到线程池执行（避免阻塞事件循环）
                result = await asyncio.get_event_loop().run_in_executor(
                    self.executor,
                    handle_task,
                    task
                )

                # 执行完成
                logger.info(f"任务 {task_id} 执行完成...")
                task.status = TaskStatus.SUCCESS
                task.title = result.title
                task.content = result.content
            except Exception as e:
                logger.exception(f"执行任务 {task_id} 发生例外 {e}")
                task.status = TaskStatus.FAILURE
                task.error_msg = str(e)

            # 保存最终结果
            async with get_session_async_ctx() as session:
                task.end_execution_time = datetime.now()
                await save_task(session, task)
        finally:
            # 删除本地文件
            Path(task.input_path).unlink()
            self.semaphore.release()
            await self.trigger_schedule()

    def close(self):
        """
        关闭
        :return:
        """
        self.closed = True
        self.executor.shutdown(wait=True)


def handle_task(task: Task) -> TaskResult:
    """
    执行转换
    :param task: 任务信息
    :return:
    """
    task_handler = task_handler_factory.get_handler(task)
    return task_handler.handle(task)


# 创建单例
task_schedule_handler = _TaskScheduleHandler(max_worker=settings.MAX_WORKER)


async def start_schedule_tasks():
    """
    启动调度任务
    """

    # 应用启动时将之前未完成的任务改为待执行状态
    async with get_session_async_ctx() as session:
        statement = select(Task).where(Task.status == TaskStatus.EXECUTING)
        result = await session.execute(statement)
        tasks = result.scalars().all()

        if tasks:
            for task in tasks:
                task.status = TaskStatus.PENDING
                task.start_execution_time = None
                task.end_execution_time = None
            session.add_all(tasks)
            await session.commit()

    asyncio.create_task(task_schedule_handler.schedule_tasks())


async def close_schedule_tasks():
    """
    关闭调度任务
    """

    task_schedule_handler.close()
